package main

import (
	"bufio"
	"context"
	"fmt"
	"github.com/volcengine/volcengine-go-sdk/service/arkruntime"
	"github.com/volcengine/volcengine-go-sdk/service/arkruntime/model"
	"github.com/volcengine/volcengine-go-sdk/service/arkruntime/utils"
	"github.com/volcengine/volcengine-go-sdk/volcengine"
	"io"
	"os"
	"strings"
	"xiaozhi/config"
)

func doubao(text string) (string, error) {
	client := arkruntime.NewClientWithApiKey(
		config.NewDouBaoLLMConfig().ApiKey,
	)

	ctx := context.Background()
	var contentBuilder strings.Builder
	fmt.Println("----- streaming request -----")
	req := model.ChatCompletionRequest{
		Model: config.NewDouBaoLLMConfig().Model,
		Messages: []*model.ChatCompletionMessage{
			{
				Role: model.ChatMessageRoleSystem,
				Content: &model.ChatCompletionMessageContent{
					StringValue: volcengine.String("你是豆包，是由字节跳动开发的 AI 人工智能助手"),
				},
			},
			{
				Role: model.ChatMessageRoleUser,
				Content: &model.ChatCompletionMessageContent{
					StringValue: volcengine.String(text),
				},
			},
		},
	}
	stream, err := client.CreateChatCompletionStream(ctx, req)
	if err != nil {
		fmt.Printf("stream chat error: %v\n", err)
		return "", err
	}
	defer func(stream *utils.ChatCompletionStreamReader) {
		err := stream.Close()
		if err != nil {
			fmt.Errorf("stream chat close error: %v\n", err)
		}
	}(stream)

	for {
		recv, err := stream.Recv()
		if err == io.EOF {
			fmt.Println("\n----- streaming end -----")
			return contentBuilder.String(), nil
		}
		if err != nil {
			fmt.Printf("Stream chat error: %v\n", err)
			return "", err
		}

		if len(recv.Choices) > 0 {
			fmt.Print(recv.Choices[0].Delta.Content)
			contentBuilder.WriteString(recv.Choices[0].Delta.Content)
		}

	}
}

func main() {
	for {

		reader := bufio.NewReader(os.Stdin) // 创建 Reader 对象
		fmt.Println("请输入内容（输入\"exit\"退出）：")

		text, err := reader.ReadString('\n') // 读取整行输入
		if err != nil {
			fmt.Println("读取输入时发生错误：", err)
			return
		}
		if text == "exit\n" {
			break
		}
		_, err = doubao(text)
		if err != nil {
			fmt.Errorf("doubao error: %v\n", err)
			return
		}
	}

}
